# xray_scan_collection.py
import numpy as np
from scipy.interpolate import griddata

from xray_scan import XRayScan


class XRayScanCollection:
    def __init__(self):
        """
        初始化XRayScanCollection对象。
        """
        self.scans = []

    def add_scan(self, xray_scan):
        """
        添加一个XRayScan对象到集合中，并保持集合按角度有序。

        :param xray_scan: 要添加的XRayScan对象
        """
        if not self.scans:
            self.scans.append(xray_scan)
        else:
            if any(len(scan.sensor_data) != len(xray_scan.sensor_data) for scan in self.scans):
                raise ValueError("All XRayScan objects must have the same sensor data length.")
            self.scans.append(xray_scan)
            self.scans.sort(key=lambda scan: scan.angle_deg)

    def get_scan_by_index(self, index):
        """
        根据索引返回按角度排序的XRayScan对象。

        :param index: 索引（0-based）
        :return: 对应的XRayScan对象
        """
        if index < 0 or index >= len(self.scans):
            raise IndexError("Index out of range.")
        return self.scans[index]

    def copy(self):
        """
        创建并返回一个XRayScanCollection对象的副本。

        :return: XRayScanCollection对象的副本
        """
        new_collection = XRayScanCollection()
        new_collection.scans = [XRayScan(scan.angle_deg, scan.sensor_data[:]) for scan in self.scans]
        return new_collection

    def __repr__(self):
        """
        返回对象的字符串表示，换行输出每个XRayScan对象。
        """
        return "XRayScanCollection:\n" + "\n".join(repr(scan) for scan in self.scans)

    def create_fourier_samples(sinogram, angles):
        """在傅里叶空间创建极坐标采样点"""
        n_angles, n_rays = sinogram.shape

        # 对每个投影进行傅里叶变换
        fourier_projections = np.fft.fft(sinogram, axis=1)
        fourier_projections = np.fft.fftshift(fourier_projections, axes=1)

        # 创建采样点坐标
        x_coords = []
        y_coords = []
        values = []

        frequencies = np.fft.fftshift(np.fft.fftfreq(n_rays))

        # print(frequencies)
        # print(n_rays)

        for i, angle in enumerate(angles):
            for j, freq in enumerate(frequencies):

                r = freq
                x = r * np.cos(angle)
                y = r * np.sin(angle)
                x_coords.append(x)
                y_coords.append(y)
                values.append(fourier_projections[i, j])

        return np.array(x_coords), np.array(y_coords), np.array(values)

    def interpolate_fourier_space(x_coords, y_coords, values, grid_size):
        """在笛卡尔网格上进行插值"""
        grid_x, grid_y = np.meshgrid(
            np.linspace(-0.5, 0.5, grid_size),
            np.linspace(-0.5, 0.5, grid_size)
        )

        # 使用griddata进行插值
        interpolated = griddata(
            (x_coords, y_coords),
            values,
            (grid_x, grid_y),
            method='nearest',
            fill_value=0
        )

        return interpolated

    def reconstruct_image(interpolated_fourier):
        """从插值后的傅里叶空间重建图像"""
        reconstructed = np.fft.ifft2(np.fft.ifftshift(interpolated_fourier))
        return np.real(reconstructed)

    def plot_sinogram(self):
        """
        绘制正弦图，展示投影数据随角度的变化
        """
        import numpy as np
        import matplotlib.pyplot as plt

        # 收集所有角度和传感器数据
        angles = [scan.angle_deg for scan in self.scans]
        sensor_data = [scan.sensor_data for scan in self.scans]

        # 转换为numpy数组以便绘图
        sinogram = np.array(sensor_data)

        # 创建图像
        plt.figure(figsize=(10, 8))
        plt.imshow(sinogram,
                   extent=[0, len(self.scans[0].sensor_data), 0, 180],
                   aspect='auto',
                   cmap='viridis')

        # 添加颜色条和标签
        plt.colorbar(label='Intensity')
        plt.xlabel('Detector Position')
        plt.ylabel('Angle (degrees)')
        plt.title('Sinogram')

        # 显示网格
        plt.grid(True, linestyle='--', alpha=0.3)

        plt.tight_layout()
        plt.show()

        return sinogram


def main():
    # 创建几个XRayScan对象
    scan1 = XRayScan(30, [0, 1, 0, 3, 0, 2, 0, 1, 0])
    scan2 = XRayScan(60, [0, 1, 0, 3, 0, 2, 0, 1, 0])
    scan3 = XRayScan(45, [0, 1, 0, 3, 0, 2, 0, 1, 0])

    # 创建XRayScanCollection对象并添加扫描
    collection = XRayScanCollection()
    collection.add_scan(scan1)
    collection.add_scan(scan2)
    collection.add_scan(scan3)

    # 打印XRayScanCollection对象
    print("Original Collection:")
    print(collection)

    # 创建XRayScanCollection对象的副本
    collection_copy = collection.copy()

    # 打印副本XRayScanCollection对象
    print("\nCopied Collection:")
    print(collection_copy)


if __name__ == "__main__":
    main()